import nibabel as nib
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import numpy as np
import os
from PIL import Image
from tqdm import tqdm
import cv2


def get_phase(s):
    global phase
    if ('-NC.' in s) or ('-Pre.' in s):
        phase = '1-NC'
    elif ('-A.' in s) or ('-Arterial.' in s):
        phase = '2-A'
    elif ('-V.' in s) or ('-Venous.' in s):
        phase = '3-V'
    elif ('-D.' in s) or ('-Delay.' in s):
        phase = '4-D'
    elif ('-T2.' in s) or ('-T2WI.' in s):
        phase = '0-T2'
    return phase


def padding_image(image_file_path, png_path):
    # 图像补全
    image = Image.open(image_file_path)
    target_size = 512
    left = (image.width - target_size) // 2  # 计算裁剪区域
    top = (image.height - target_size) // 2
    right = left + target_size
    bottom = top + target_size

    cropped_image = image.crop((left, top, right, bottom))  # 裁剪图像
    # print("补全前后对比:", image.size, cropped_image.size)
    cropped_image.save(png_path)  # 保存
    image.close()


def convert_3d_nii_to_png(nii_path, output_png_folder, label=False):
    """将3D的nii文件转换为身体横截方向(tra)的png序列.

    :param nii_path: 源nii文件路径.
    :param output_png_folder: 输出的png的文件夹路径.
    """
    if output_png_folder != "" and not os.path.exists(output_png_folder):
        os.mkdir(output_png_folder)  # 如果保存文件夹不存在，则创建'

    nii_img = nib.load(nii_path)  # 加载NIfTI文件
    nii_data = nii_img.get_fdata()  # 获取NIfTI数据

    s = nii_img.get_filename()
    nii_name = s[s.rfind('/') + 1:]
    phase = get_phase(s)  # 获取期相 (NC/A/V/D)

    first_dash_index = nii_name.find('-')
    second_dash_index = nii_name.find('-', first_dash_index + 1)
    patient_name = nii_name[first_dash_index + 1:second_dash_index + 1]  # MOU MOU MOU-
    if label:  # 控制data和label的文件名不同
        png_name = patient_name + phase
    else:
        png_name = patient_name + phase

    nii_z_num = nii_data.shape[2]  # 遍历数据中的每个切片
    for i in tqdm(range(nii_z_num), ncols=50):
        slice_data = nii_data[:, :, i]
        rotated_slice = np.rot90(slice_data, k=1)  # 顺时针旋转90度
        mirrored_slice = np.flipud(rotated_slice)  # 上下镜像

        # 转灰度图
        # image = Image.fromarray(slice_data)
        # if image.mode != 'L':
        #     image = image.convert('L')
        # temp1_file_path = os.path.join(output_png_folder, f'{output_png_folder}/{png_name}-{i:03d}.png')
        # image.save(temp1_file_path)
        # image.close()

        plt.imshow(mirrored_slice, cmap='gray')
        plt.axis('off')
        # 保存图像为 PNG 文件
        # temp1_file_path = os.path.join(output_png_folder, "temp.png")
        temp1_file_path = os.path.join(output_png_folder, f'{output_png_folder}/{png_name}-{(nii_z_num - i):03d}.png')
        plt.savefig(temp1_file_path, bbox_inches='tight', pad_inches=0, dpi=86.65)
        plt.close()
        # break

        # if label:  # label不需要调节对比对
        #     png_path = os.path.join(output_png_folder, f'{output_png_folder}/{png_name}-{i:03d}.png')
        #     padding_image(temp1_file_path, png_path)
        #     if os.path.exists(temp1_file_path):
        #         os.remove(temp1_file_path)  # 删除临时文件
        # else:
        #     temp2_file_path = os.path.join(output_png_folder, "temp_crop.png")
        #     padding_image(temp1_file_path, png_path)
        #
        #     image = cv2.imread(temp2_file_path, cv2.IMREAD_GRAYSCALE)
        #     clip_limit = 2.0  # 对比度
        #     clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
        #     clamped_image = clahe.apply(image)
        #     cv2.imwrite(f'{output_png_folder}/{png_name}-{i:03d}.png', clamped_image)
        #
        #     if os.path.exists(temp2_file_path):
        #         os.remove(temp1_file_path)  # 删除临时文件
        #         os.remove(temp2_file_path)


if __name__ == '__main__':
    all_path = r"C:\Users\18049\Desktop\nii"
    dcm_path_2022 = all_path + r"\1_unclassified_dcm\metastasis 2022"
    # nii_path_2022 = all_path + r"\2_classification_nii\metastasis 2022"
    nii_path_2022_resized = all_path + r"\2_classification_nii\metastasis 2022 resized"

    nii_path = r'C:\Users\18049\Desktop\nii\3_mark_label\1-Pre'
    output_data_folder = r"C:\develop\GraduationDesign\unet_seg\data\train\Data"
    output_label_folder = r"C:\develop\GraduationDesign\unet_seg\data\train\Ground"
    namelist = ['1', '2', '3', '4', '5', '6']
    i = 0
    for file in os.listdir(nii_path):  # 遍历图像列表
        if file.endswith(".nii"):
            # NIfTI文件路径
            file_path = os.path.join(nii_path, file)
            print("正在处理:", file_path)
            if 'tumor-' in file:
                convert_3d_nii_to_png(file_path, output_label_folder + namelist[i], label=True)  # 目标标签
            else:
                convert_3d_nii_to_png(file_path, output_data_folder + namelist[i], label=False)  # 图像数据
            print('-' * 50, sep='')
            i = i + 1



    # equal_nii = r'C:\develop\GraduationDesign\medical-image-preprocessing\test\equalized_image1-1-Pre.nii'
    # equal_nii = r"C:\Users\18049\Desktop\nii\2_classification_nii\metastasis 2022 resized\WANG CUI HUA_M00263248\WANG CUI HUA-M00263248-1-Pre.nii"
    # convert_3d_nii_to_png(equal_nii, output_data_folder, label=False)

    # for patient_dcm_folder in os.listdir(nii_path_2022_resized):
    #     patient_dcm_folder_path = os.path.join(nii_path_2022_resized, patient_dcm_folder)  # C:\..\CHEN YI JIANG_M00xxxxx
    #     if os.path.isdir(patient_dcm_folder_path):  # 检查是否是文件夹
    #         print('-' * 50, '\n', patient_dcm_folder, sep='')
    #         patient_nii_folder_path = os.path.join(nii_path_2022_resized, patient_dcm_folder)  # 创建存放目标nii的文件夹
    #
    #         for file in os.listdir(patient_nii_folder_path):  # 遍历图像列表
    #             if file.endswith(".nii"):
    #                 # NIfTI文件路径
    #                 file_path = os.path.join(patient_nii_folder_path, file)
    #                 print("正在处理:", file_path)
    #                 if 'tumor-' in file:
    #                     convert_3d_nii_to_png(file_path, output_label_folder, label=True)  # 目标标签
    #                 else:
    #                     convert_3d_nii_to_png(file_path, output_data_folder, label=False)  # 图像数据
    #                 print('-' * 50, sep='')

